import numpy as np
import os
import output_function as of
import pandas as pd
from scipy.sparse import csr_matrix

def vel_to_omega(u, v, dx, dy):  
    omega = (v[1:-1, 2:] - v[1:-1, :-2]) / 2 / dx - (u[2:, 1:-1] - u[:-2, 1:-1]) / 2 / dy  
    return omega

#扩散算符
def diff_operator(myfun,nu,dx,dy):
    diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
    diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
    return diff_x+diff_y

def laplace(n_x, n_y, dx, dy):
    # 按公式敲依次敲出对角阵/移格对角阵，然后组合
    Dx = np.diag(np.ones(n_x)) * 2 - np.diag(np.ones(n_x - 1),1) - np.diag(np.ones(n_x - 1), -1)
    # Dx[0, 0] = 1
    # Dx[-1, -1] = 1
    Ix = np.diag(np.ones(n_y))
    Dy = np.diag(np.ones(n_y)) * 2 - np.diag(np.ones(n_y - 1),1) - np.diag(np.ones(n_y - 1), -1)
    # Dy[0, 0] = 1
    # Dy[-1, -1] = 1
    Iy = np.diag(np.ones(n_x))
    L = np.kron(Ix, Dx) / dx ** 2 + np.kron(Dy, Iy) / dy ** 2       # 克罗内克积
    L_sparse = csr_matrix(L)
    return L_sparse

'''计算流函数，本质上是解泊松方程。要注意边界条件的处理'''
def compute_stream(omega_int, dx, dy, tol=1e-7, psi_init=None,max_iters=100000):
    """
    Jacobi iteration
    主输入:涡度场值omega_int
    辅输入:dx,dy,网格长,用于算差分矩阵
    主返回:计算得到的psi_init
    """
    if psi_init is None:
        psi_init = np.zeros((Ny, Nx))
    change = 1.
    iter = 0
    while (change > tol and iter<max_iters) or iter<1000:
        psi_new = psi_init.copy()
        #对中间区域格点的迭代过程，这里把差分矩阵在雅克比迭代情况的表达式写进来了。计算采用矩阵加法）
        psi_new[1:-1, 1:-1] = (omega_int \
                               + (psi_init[2:, 1:-1] + psi_init[:-2, 1:-1]) / dy / dy \
                               + (psi_init[1:-1, 2:] + psi_init[1:-1, :-2]) / dx / dx) / (2 / dx / dx + 2 / dy / dy)
        '''
        如果边界上流函数一直为零,初始时设成0就好,后面不动它
        顶盖第二类边条件(收敛性存在问题):psi_new[-1, 1:-1]=psi_new[-2, 1:-1]+dy
        '''
        #psi_new[-1, :]=psi_new[-2, :]   #顶部
        psi_new[-1, :]=0
        psi_new[0, :]=0  #底部
        psi_new[1:-1,0]=0
        psi_new[1:-1,-1]=0
        
        change = np.max(np.abs(psi_new - psi_init))  #change为迭代后矩阵元素变化最大值

        psi_init = psi_new
        iter += 1   
    end_signal = True if iter == 1 else False  #只用迭代1次时，end_signal返回真，这将用于停止迭代。因此此程序似乎在求从零直到稳态的过程。
    if iter>max_iters/10:
        print(iter)

    return psi_init, end_signal

# 使用示例
base_path = './data/RB_cavity/'  # 根据需要修改为您的保存路径
out_iter =20000  # 假设我们要加载第300时间步的数据

# 加载数据
n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')

n_density=n

# 加载模拟参数
config = of.load_simulation_config(f'{base_path}simulation_config.json')
# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Nx = config['Nx']
Ny = config['Ny']
n_up = config['n_up']
n_0 = config['n_0']
n_down = config['n_down']
Delta_n = config['Delta_n']
Ra=Ra_star = config['Ra_star']
prandtl=Pr = config['Pr']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']

# 打印一些变量以确认它们已被正确加载和赋值
print(f"dx: {dx}")
print(f"Ra_star: {Ra_star}")
print(f"dt: {dt}")
# 确认是否成功加载了正确的时间步长
if out_iter_loaded == out_iter:
    print(f"Successfully loaded data for time step: {out_iter_loaded}")
    print(omega.shape)
else:
    print("Failed to load the correct time step or file does not exist.")

# 接下来，您可以使用加载的数据进行分析、绘图或其他处理
print(xx)
# error=vel_to_omega(u,v,dx,dy)-omega[1:-1,1:-1]
# print(error)

#psi, end_signal = compute_stream(omega[1:-1, 1:-1], dx, dy, tol=1e-7, psi_init=psi)

L = laplace(Nx, Ny, dx, dy)
print("L矩阵的维度是:", L.shape)

# error=diff_operator(psi,1,dx,dy)+omega[1:-1,1:-1]
# print(error)


# data = pd.DataFrame(error)
# writer = pd.ExcelWriter(f'{base_path}omega_error_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()






    
# data = pd.DataFrame(omega)
# writer = pd.ExcelWriter(f'{base_path}omega_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()

# data = pd.DataFrame(u)
# writer = pd.ExcelWriter(f'{base_path}u_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()

# data = pd.DataFrame(v)
# writer = pd.ExcelWriter(f'{base_path}v_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()
    
# omega_drive_term=-Ra*prandtl*(n_density[1:-1,2:]-n_density[1:-1,:-2])/2/dx
# #print(omega_drive_term)
# data = pd.DataFrame(omega_drive_term)
# writer = pd.ExcelWriter(f'{base_path}omega_drive_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()

# def diff_operator(myfun,nu):
#     diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
#     diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
#     return diff_x+diff_y

# omega_diff=diff_operator(omega,Pr)
# data = pd.DataFrame(omega_diff)
# writer = pd.ExcelWriter(f'{base_path}omega_diff_{out_iter}.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()

# 使用加载的数据进行分析、绘图或其他处理
# base_folder=f'./RB_cavity_plot_Ra_{Ra_star}_Lx_{Lx}'
# of.diag_n_pcolor_filepath(n, n_init0_matrix, u, v, xx, yy, out_iter, Ra_star, dt,base_folder)
# of.plot_and_save_n_density_x_section(x,y, n, n_init0_matrix, Ny//2, out_iter,title_base="delta n_density Distribution",save_folder_base=base_folder)
# of.plot_and_save_v_x_section(x,y, v, Ny//2, out_iter,title_base="V Distribution along X-axis",save_folder_base=base_folder)